
import sys

def geneSummary(filename, norm=None):
  READ_LENGTH = 32
  virtualLength = {}
  totalLength = {}
  totalExpression = {}
  totalSquaredExpression = {}
  chr = {}
  strand = {}
  minPos = {}
  maxPos = {}
  intronStats = {}
  exonStats = {}
  geneStats = {}
  for line in open(filename):
    if line[0]=="#": continue
    cols = line.strip().split("\t")
    (gene, exonOrIntron, exonOrIntronNum, exonOrIntronId, geneChr, geneStrand, rangeMin, rangeMax, totalCov, totalSquaredCov) = tuple(cols)
    rangeMin = long(rangeMin)
    rangeMax = long(rangeMax)
    if exonOrIntron=="Exon":
      if not totalExpression.has_key(gene):
        totalExpression[gene] = 0.0
        totalSquaredExpression[gene] = 0.0
        totalLength[gene] = 0
        virtualLength[gene] = 0
        chr[gene] = geneChr
        strand[gene] = geneStrand
        minPos[gene] = rangeMin
        maxPos[gene] = rangeMax
      maxPos[gene] = max(maxPos[gene], rangeMax)
      exonLength = rangeMax - rangeMin + 1
      virtualExonLength = 0
      if exonLength >= READ_LENGTH:
        virtualExonLength = exonLength - READ_LENGTH + 1
      sm = float(totalCov)
      smsq = float(totalSquaredCov)
      participationRatio = 0.0
      if sm>0.0: participationRatio = (sm * sm) / (exonLength * smsq)
      aveExpression = sm / (1.0 * exonLength)
      # print "%s\t%f\t%f" % (line.strip(), aveExpression, participationRatio)
      exonStats[(gene, exonOrIntronNum)] = [aveExpression, participationRatio]
      virtualLength[gene] += virtualExonLength
      totalLength[gene] += exonLength
      totalExpression[gene] += sm
      totalSquaredExpression[gene] += smsq
    elif exonOrIntron=="Intron":
      intronLength = rangeMax - rangeMin + 1
      sm = float(totalCov)
      smsq = float(totalSquaredCov)
      participationRatio = 0.0
      if sm>0.0: participationRatio = (sm * sm) / (intronLength * smsq)
      aveExpression = sm / (1.0 * intronLength)
      # print "%s\t%f\t%f" % (line.strip(), aveExpression, participationRatio)
      intronStats[(gene, exonOrIntronNum)] = [aveExpression, participationRatio]
  for gene in totalExpression.keys():
    sm = totalExpression[gene]
    smsq = totalSquaredExpression[gene]
    participationRatio = 0.0
    if sm>0.0: participationRatio = (sm * sm) / (totalLength[gene] * smsq)
    aveExpression = 0.0

    # Correct virtual length... splice junctions are included, so only
    # positions near the ends, not near all junctions, need to be excluded.
    virtualLength[gene] = totalLength[gene] - READ_LENGTH + 1

    if virtualLength[gene]>0.0: aveExpression = sm / (1.0 * virtualLength[gene])
    geneStats[gene] = [aveExpression, participationRatio]
    if norm==None:
      aveExpression = "%f" % aveExpression
    else:
      aveExpression = "%f\t%f" % (aveExpression, aveExpression * norm)
    print "%s\tGeneExpression\t%s\t%s\t%s\t%ld\t%ld\t%d\t%d\t%f\t%f\t%s\t%f" % (
      gene, "%s_%s_%s_%s" % (chr[gene], strand[gene], minPos[gene], maxPos[gene]),
      chr[gene], strand[gene], minPos[gene], maxPos[gene], totalLength[gene], virtualLength[gene],
      sm, smsq, aveExpression, participationRatio)
  for line in open(filename):
    if line[0]=="#": continue
    cols = line.strip().split("\t")
    gene = cols[0]
    geneExpression = geneStats[gene][0]
    if cols[1]=="Exon":
      es = exonStats[(cols[0], cols[2])]
      expr = es[0]
      pratio = es[1]
      normExpr = expr / 1.0e-6
      if geneExpression > 0.0: normExpr = expr / geneExpression
      if norm==None:
        expr = "%f" % expr
      else:
        expr = "%f\t%f" % (expr, expr * norm)
      print "%s\t%s\t%f\t%f" % (line.strip(), expr, normExpr, pratio)
    elif cols[1]=="Intron":
      ins = intronStats[(cols[0], cols[2])]
      expr = ins[0]
      pratio = ins[1]
      normExpr = expr / 1.0e-6
      if geneExpression > 0.0: normExpr = expr / geneExpression
      if norm==None:
        expr = "%f" % expr
      else:
        expr = "%f\t%f" % (expr, expr * norm)
      print "%s\t%s\t%f\t%f" % (line.strip(), expr, normExpr, pratio)

def printUsage(args, usage, minArgs=2):
  if len(args) < minArgs or args[1].startswith("-h") or args[1].startswith("--h"):
    print usage
    sys.exit(1)

usageString = """
Usage: python %s EXON_INTRON_FILE [NORM]

Calculates the expression level and participation ratio of each exon,
intron, and gene in the specified exon/intron coverage file (as produced
by 'calculateGeneCoverage.py').  The exon/intron rows in the output have
the same columns as in the exon/intron coverage file, with the following
columns appended:

11. Intron or exon expression (total coverage divided by length)
(11N. Intron or exon expression times NORM, if NORM was specified.)
12. Intron or exon expression divided by gene expression
13. Intron or exon participation ratio

The gene expression rows have the following columns:

1. Gene name
2. 'GeneExpression'
3. Gene position range ('chr1_+_24722_29182')
4. Chromosome ('chr12')
5. Strand ('+' / '-')
6. First position
7. Last position
8. Total length (sum of exon lengths)
9. Virtual length (# of possible alignment positions: sum of exon lengths, minus 31)
10. Total coverage (sum of total coverage of all exons)
11. Total squared coverage (sum of total squared coverage of all exons)
12. Gene expression level (total coverage divided by virtual length)
(12N. Gene expression times NORM, if NORM was specified.)
13. Gene participation ratio
"""
    
if __name__=="__main__":
  printUsage(sys.argv, usageString % sys.argv[0])
  if len(sys.argv)==2:
    geneSummary(sys.argv[1])
  else:
    geneSummary(sys.argv[1], float(sys.argv[2]))

